Description

This PR optimizes grouped MXFP8 quantization by introducing layout-specific kernel configurations and work scheduling strategies.

Groups with uniform shapes or a common last dimension use direct CTA-to-chunk mapping, while layouts requiring per-tensor TMA descriptors use per-tensor persistent scheduling. The kernel processes larger chunks through multiple stages, reducing scheduling and tensor-mapping overhead for large grouped workloads.

The change also ensures that dynamically updated TMA descriptors are acquired before use, handles graph-safe logical capacity without out-of-bounds accesses, and preserves correct dBias workspace behavior across multi-chunk execution.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

• Extend grouped TMA storage with per-tensor row, column, and offset metadata.
• Introduce layout-specific grouped MXFP8 kernel configurations and compile-time traits.
• Use direct CTA-to-chunk mapping for SAME_BOTH_DIMS and VARYING_FIRST_DIM layouts.
• Use per-tensor persistent scheduling for VARYING_LAST_DIM and VARYING_BOTH_DIMS layouts.
• Process 128x128 and 128x256 chunks through multiple X/Y stages.
• Acquire dynamically updated TMA descriptors before they are consumed by the quantization kernel.
• Avoid out-of-bounds accesses when graph-safe logical capacity exceeds the active tensor range.
• Preserve correct rowwise, columnwise, and bidirectional scaling behavior.
• Preserve correct dBias accumulation and workspace reduction for multi-chunk execution.
• Add focused regression coverage for all grouped shape representations, multi-chunk dBias/dAct execution, and the maximum of 64 dynamic TMA descriptors.
• Keep all existing dispatch paths, operator tests, and CUDA architecture targets enabled.

Runtime Speedup Factor on GB300

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes